import joblib
import pandas as pd
import numpy as np
from sklearn.feature_extraction import DictVectorizer
from sklearn.model_selection import  train_test_split
from sklearn.tree import DecisionTreeClassifier, export_graphviz

titanicTrain = pd.read_csv("./Data/titanic_train.csv")
#titanicTest = pd.read_csv("./Data/titanic_test.csv")
print("titanicTrain:")
print(titanicTrain)
x = titanicTrain[["Pclass", "Age", "Sex"]]
y = titanicTrain["Survived"]
x["Age"].fillna(x["Age"].mean(), inplace=True)
#x["Age"].fillna(x["Age"].mean(), inplace=True)
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.25)


#特征值提取，这里用one-hot representation
x_train = x_train.to_dict(orient="records")
x_test = x_test.to_dict(orient="records")

print("x_train:")
print(x_train)
print("x_test:")
print(x_test)

transfer = DictVectorizer(sparse=False)

x_train = transfer.fit_transform(x_train)
print("转换为count-based后的x_train为：")
print(x_train)
x_test = transfer.transform(x_test)
print("转换为count-based后的x_test为：")
print(x_test)
#决策树
estimator = DecisionTreeClassifier(max_depth=3)
estimator.fit(x_train, y_train)

y_pre = estimator.predict(x_test)
print("预测结果：", y_pre)
print("真实结果：", y_test)
print("预测的准确率：")
print(estimator.score(x_test, y_test))

# 导出决策树结构 feature_names根据特征值 dict.get_feature_names()
'''
dot_data = export_graphviz(estimator, out_file="./T_Tree.dot",
                           feature_names=['年龄', 'pclass=1st', 'pclass=2nd', 'pclass=3rd', '女性', '男性'])
'''
# 保存模型
#joblib.dump(estimator, 'ttnk.joblib')

